"""
a script to split a data dir into trainset and valset, with all files that ends with jpg or jpeg, store their filenames and split with an factor
"""
import os
import random
from shutil import copyfile
import warnings
def split_data(data_dir, val_size):

    print("spliting data")
    files = os.listdir(data_dir)
    files = [file for file in files if file.endswith(".jpg") or file.endswith(".jpeg")]
    random.shuffle(files)

    if val_size == 0:
        # copy train.txt to val.txt
        train_files = files
        
        val_files = train_files
        with open(os.path.join('splits',"train.txt"), "w") as f:
            for file in train_files:
                f.write(file + "\n")    
        warnings.warn('did not split data, using the whole trainset as valset')
        copyfile(os.path.join('splits',"train.txt"), os.path.join('splits',"val.txt"))
        # return

    else:
        
        split_point = int(len(files) * val_size)
        train_files = files[:-split_point]
        val_files = files[-split_point:]

        # then write the trainfile names in to train.txt, the other to val.txt
        with open(os.path.join('splits',"train.txt"), "w") as f:
            for file in train_files:
                f.write(file + "\n")

        # else:
        with open(os.path.join('splits',"val.txt"), "w") as f:
            for file in val_files:
                f.write(file + "\n")

    print(f'{len(files)=}')
    print(f'{len(train_files)=}')
    print(f'{len(val_files)=}')

  
if __name__ == "__main__":
    random.seed(42)
    if os.getenv('face_dev'):
        data_dir = "../data/2792"
    else:
        data_dir = '/home/data/2792'
    if os.path.exists(os.path.join(data_dir,"train.txt")) and os.path.exists(os.path.join(data_dir,"val.txt")):
        # print("train.txt or val.txt already exists, please remove them first")
        # exit()
        print("exist train.txt, will overwrite", )
    else:
        print("train.txt and val.txt will be created")
        
    val_size = 0.1
    split_data(data_dir, val_size)
